﻿#pragma once

#include <tuple>
#include <vector>

#include <xy/math/matrix.h>

namespace xy
{

namespace solver
{

struct line_search
{
    /**
     * @brief 进退法搜索
     *
     * @tparam T
     * @param[in] f
     * @param[in] x 初始位置
     * @param[in] d 搜索方向
     * @param[in] h 步长
     * @return std::tuple<double, double>
     */
    template <class T>
    [[nodiscard]] std::tuple<double, double> static forward_backward_line_search(double (*f)(T), T x, T d, double h)
    {
        int t = 2;

        double a = 0;
        double b = h;
        double fx0 = f(x);

        int k = 0;
        while (k < 30)
        {
            double fx1 = f(x + d * b);
            if (fx1 < fx0)
                b *= t;
            else
            {
                if (k != 0)
                    break;

                // 如果第一步上升，则记录上升点，反向搜索高点
                a = b;
                b = -h;
            }

            fx0 = fx1;
            k++;
        }

        return {std::min(a, b), std::max(a, b)};
    }

    /**
     * @brief 黄金分割法搜索
     *
     * @tparam T
     * @param[in] f
     * @param[in] x 初始位置
     * @param[in] d 搜索方向
     * @param[in] tol
     * @return double
     */
    template <class T> [[nodiscard]] static double golden_section_line_search(double (*f)(T), T x, T d, double tol)
    {
        static const double rb = (std::sqrt(5.0) - 1) / 2;
        static const double ra = 1 - rb;

        auto [a, b] = forward_backward_line_search<T>(f, x, d, 1e-2);

        // 选取试探点
        double lambda = a + ra * (b - a);
        double mu = a + rb * (b - a);

        // 计算试探点的值
        std::vector<double> vals = {a, lambda, mu, b};
        std::vector<double> fs = {f(x + d * a), f(x + d * lambda), f(x + d * mu), f(x + d * b)};

        // 同时比较 a, b, lambda, mu 处的值
        std::size_t t = std::distance(fs.begin(), std::min_element(fs.begin(), fs.end()));

        // 迭代
        while (vals[3] - vals[0] > tol)
        {
            if (t < 2)
            {
                vals[3] = vals[2];
                vals[2] = vals[1];
                vals[1] = vals[0] + ra * (vals[3] - vals[0]);
                fs[3] = fs[2];
                fs[2] = fs[1];
                fs[1] = f(x + d * vals[1]);

                if (fs[t + 1] < fs[t])
                    t++;
            }
            else
            {
                vals[0] = vals[1];
                vals[1] = vals[2];
                vals[2] = vals[0] + rb * (vals[3] - vals[0]);
                fs[0] = fs[1];
                fs[1] = fs[2];
                fs[2] = f(x + d * vals[2]);

                if (fs[t - 1] < fs[t])
                    t--;
            }
        }

        return (vals[0] + vals[3]) / 2;
    }

    /**
     * @brief Armijo-Goldstein 线性搜索
     *
     * @tparam T
     * @param[in] f
     * @param[in] df
     * @param[in] x 初始位置
     * @param[in] d 下降方向
     * @param[in] tol
     * @return double
     */
    template <class T>
    [[nodiscard]] static double armijo_goldstein_line_search(double (*f)(T), T (*df)(T), T x, T d, double tol)
    {
        double a = 0, b = std::numeric_limits<double>::infinity();
        double alpha = 10;
        double rho = 0.25;
        double t = 2;
        double fx = f(x);
        T dfx = df(x);

        while (b - a > tol)
        {
            double falpha = f(x + d * alpha);

            // 右侧判断，满足则更新左端；否则更新右端
            if (falpha < fx + alpha * rho * dot(dfx, d))
            {
                // 左侧判断，满足则直接退出
                if (falpha > fx + alpha * (1 - rho) * dot(dfx, d))
                    break;

                // 否则增大左端
                a = alpha;

                // 如果右端无穷大，则可以增大试探步长。否则减小右端
                if (b == std::numeric_limits<double>::infinity())
                {
                    alpha *= t;
                    continue;
                }
            }
            else
                b = alpha;

            // 调整试探步长，使其靠近 [a, b] 的中点
            alpha = (a + b) / 2;
        }

        return alpha;
    }

    /**
     * @brief Wolfe-Powell 线性搜索
     *
     * @tparam T
     * @param[in] f
     * @param[in] df
     * @param[in] x
     * @param[in] d
     * @param[in] tol
     * @return double
     */
    template <class T>
    [[nodiscard]] static double wolfe_powell_line_search(double (*f)(T), T (*df)(T), T x, T d, double tol)
    {
        double a = 0, b = std::numeric_limits<double>::infinity();
        double alpha = 10;
        double rho = 0.25;
        double sigma = 0.5;
        double t = 2;
        double fx = f(x);
        T dfx = df(x);

        while (b - a > tol)
        {
            double falpha = f(x + d * alpha);
            T dfalpha = df(x + d * alpha);

            // 右侧判断，满足则更新左端；否则更新右端
            if (falpha < fx + alpha * rho * dot(dfx, d))
            {
                if (dot(dfalpha, d) > sigma * dot(dfx, d))
                    break;

                // 否则增大左端
                a = alpha;

                // 如果右端无穷大，则可以增大试探步长。否则减小右端
                if (b == std::numeric_limits<double>::infinity())
                {
                    alpha *= t;
                    continue;
                }
            }
            else
                b = alpha;

            // 调整试探步长，使其靠近 [a, b] 的中点
            alpha = (a + b) / 2;
        }

        return alpha;
    }

    /**
     * @brief 回溯线性搜索
     *
     * @tparam T
     * @param[in] f
     * @param[in] df
     * @param[in] x
     * @param[in] d
     * @param[in] tol
     * @return double
     */
    template <class T>
    [[nodiscard]] static double backtracking_line_search(double (*f)(T), T (*df)(T), T x, T d, double tol)
    {
        double a = 0, b = std::numeric_limits<double>::infinity();
        double alpha = 10;
        double rho = 0.25;
        double sigma = 0.5;
        double t = 2;
        double fx = f(x);
        T dfx = df(x);

        while (b - a > tol)
        {
            double falpha = f(x + d * alpha);
            T dfalpha = df(x + d * alpha);

            // 右侧判断，满足则更新左端；否则更新右端
            if (falpha < fx + alpha * rho * dot(dfx, d))
                break;

            b = alpha;
            alpha = -dot(dfx, d) * alpha * alpha / (2 * (falpha - fx - dot(dfx, d) * alpha));
        }

        return alpha;
    }
};

struct optimal
{
    /**
     * @brief 最速下降法
     *
     * @param[in] f
     * @param[in] df
     * @param[in] x
     * @param[in] tol
     * @return T
     */
    template <class T> [[nodiscard]] static T steepest_descent(double (*f)(T), T (*df)(T), T x, double tol)
    {
        T d = -df(x);
        while (d.norm() > tol)
        {
            double alpha = line_search::golden_section_line_search(f, x, d, tol);
            x += d * alpha;
            d = -df(x);
        }
        return x;
    }

    /**
     * @brief 纯牛顿法
     *
     * @tparam T
     * @tparam U
     * @param[in] f
     * @param[in] hf
     * @param[in] x
     * @param[in] tol
     * @return T
     */
    template <class T, class U = T> [[nodiscard]] static T pure_newton(T (*df)(T), U (*hf)(T), T x, double tol)
    {
        T g = df(x);
        while (g.norm() > tol)
        {
            x -= g / hf(x);
            g = df(x);
        }
        return x;
    }

    /**
     * @brief 线搜索牛顿法
     *
     * @tparam T
     * @tparam U
     * @param[in] f
     * @param[in] df
     * @param[in] hf
     * @param[in] x
     * @param[in] tol
     * @return T
     */
    template <class T, class U = T>
    [[nodiscard]] static T line_search_newton(double (*f)(T), T (*df)(T), U (*hf)(T), T x, double tol)
    {
        T g = df(x);
        while (g.norm() > tol)
        {
            auto d = -g / hf(x);
            x += d * line_search::golden_section_line_search<T>(f, x, d, tol);
            g = df(x);
        }
        return x;
    }
};

} // namespace solver

} // namespace xy
